[OMNIML-6060] Load an expert-parallel plan through load_pretrained - #2780

Claude review passed — no blocking issues found. LGTM

0 CRITICAL, 0 IMPORTANT, 0 new SUGGESTION. Full review; the trigger was a bare /claude review, so no scoping applied. All 6 changed files reviewed (+249/-58, 92 source insertions): modelopt/torch/utils/plugins/model_load_utils.py and CHANGELOG.rst first, then all four test files. Read _resolve_checkpoint_dir, _route, _load_sharded, load_pretrained's option plumbing, _decide_distribution, _dense_from_env, _ep_degree, resolve_distribution_plan, load_and_prepare_ep (config read, _ckpt_quant_method gates, from_pretrained_kwargs) and print_rank_0's import for the context the diff lacks. I could not execute the suites in this sandbox, so the test reasoning below is by reading, not by running.

Round-3 fixes verified

HEAD is d2db43ff; the delta since the last approval (ae1af9af) is exactly the two SUGGESTIONs that round left open, both now closed:

1. Hub options reach the weight download. _resolve_checkpoint_dir(ckpt_path, rank, **hub_options) forwards to both snapshot_download calls, so the rank-0 prefetch and the post-barrier resolve agree. hub_options is built in load_pretrained from exactly {"token", "cache_dir", "local_files_only"} (line 676) — all three are real snapshot_download parameters, and revision is deliberately kept out in its own dict for ensure_local_checkpoint, so no unexpected keyword can reach the download. A direct caller of the public load_and_prepare_ep handing it a Hub ID now gets token honored by the download, one cache_dir for config and weights, and local_files_only=True respected rather than silently fetching. load_and_prepare_ep copies with dict(hub_options or {}), so the caller's dict is not mutated. The fix has a test that fails without it: test_the_weights_load_in_the_checkpoints_dtype_through_its_class now pins downloads == [{"local_files_only": True}] alongside the config read and from_pretrained.
2. The replicated-MoE note goes through print_rank_0 (imported at line 55), so it no longer depends on the caller invariant that _decide_distribution runs only in broadcast_from's rank-0 callable. Cosmetic; no test added, which matches the repo's "be stingy" guidance for a shared helper already covered by test_an_ep_capable_moe_is_planned_expert_parallel's rank assertion.

Re-derived rather than taken on trust

• Every env-knob path still terminates. I rewalked the knob combinations an env-driven torchrun can reach on an EP-capable MoE: FSDP_SIZE=1 alone raises naming EP_DEGREE=1, then loads; EP_DEGREE=1 alone either raises (small model, dense resolves to replicate) or FSDP2-shards (big model, which then refuses a quantized checkpoint), then loads with FSDP_SIZE=1; EP_DEGREE=1 FSDP_SIZE=world FSDP2-shards; both at 1 replicates with the note. asked_to_replicate is latched after _ep_degree/_dense_from_env resolve but before "auto" becomes concrete, and cannot go stale — when it is True, none of the ep_size == "auto"/> 1 or dense == "auto"/"fsdp" branches can fire, so both values are unchanged at the check.
• math.gcd(num_experts, world_size) == 1 is exactly the right test, not merely close: any common divisor d > 1 divides world_size (so DistributionPlan.__post_init__ passes) and num_experts (so the loader's num_experts % ep_degree check passes), so gcd > 1 guarantees a usable degree and gcd == 1 rules every one out. The third branch's advice stays reachable because _choose_ep_degree returns the largest common divisor, so gcd > 1 with ep_size == 1 can only come from an explicit 1. The not num_experts branch closes the earlier num_experts=0 interpolation.

What checked out

• Routing. _route returns _EXPERT_PARALLEL ahead of the is_fsdp_dense test, so ep_size>1, dense="fsdp" reaches load_and_prepare_ep, which shards the dense part over the ep group itself — the right owner. plan threads positionally to match the signature. _SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL] is exactly the loader's remaining parameters and rightly omits cpu_offload/mp_policy/broadcast_chunk_size; hub_options is deliberately not in it, so the explicit keyword cannot collide with a caller-supplied **kwargs entry. All five names have real runtime use paths — no declaration-only arguments.
• Refusal placement. The dispatch sits after the shared auto_class, dtype, unknown-kwarg and checkpoint_has_mtp refusals (so the Nemotron-H MTP claim holds) and before the FSDP2-only quantization_config/_auto_class checks — correct, since the loader has its own _EP_SOURCE_FORMATS/fp8-per-tensor/fp8-unverified gates and its own class pick. A hand-written DistributionPlan(8, ep_size=8, dense="fsdp") skips the decision's unshardable refusals, but the loader repeats each before loading (is_fsdp_dense and edp > 1, is_fsdp_dense on block-fp8, quant method, empty/unstyled base_model_ep_plan, num_experts % ep_degree).
• Composed-path options. With a decided plan kwargs.setdefault("hf_config", ...) makes the loader skip its config read; attn_implementation is still applied (it goes into from_pretrained_kwargs independently of hf_config) and the handed config really drives the build (from_pretrained_kwargs["config"]) rather than being read only for quant_method/num_experts/ep_plan. With a handed plan hf_config stays None, so the loader's own config read runs — and that is the read the new hub_options now reaches.
• Distributed safety. The RuntimeError fires inside broadcast_from("plan", ...)'s rank-0 callable, which pickles and re-raises on every rank, so the launch fails rather than stranding ranks 1..N. No new rank-dependent collective. _route's plan.world_size != dist.size() check means an EP plan cannot reach the loader without a live group — why the EP branch needs no dist.setup() of its own. The world_size == 1 early return precedes the new guard, so the "load it in a single process" remedy really works.
• Record round-trip. The loader calls record_unplaced_source_keys before returning, so _modelopt_unplaced_source_keys exists when load_pretrained reads it; it records no plan, so the caller's single set_distribution_plan does not trip the already-recorded guard.
• No in-repo flow regresses. examples/hf_ptq/hf_ptq.py and example_utils.py are the only callers; neither passes distribution= nor uses distributed_run, so both stay on the world_size == 1 early return and never reach the refusal.
• Tests. The dispatch test is mutation-sensitive as claimed (loaders == {}, auto_class.calls == []) and pins hub_options == {"local_files_only": True}; the base_model_ep_plan = None DeepSeek-V3 config keeps it off the transformers-5.7 skip, sound because an explicit distribution= bypasses the decision. Parametrizing test_a_sharded_load_refuses_what_it_cannot_honor over both shardings extends rather than adds, and the deleted test_no_ep_degree_dividing_the_experts_leaves_them_whole is replaced by the auto_finds_no_degree parametrization of the inverted behavior rather than just dropped. In the GPU test local.shape[0] == per_rank plus the owned contiguous-slice compare is what makes the EP case honest — not sharded - experts alone would pass if nothing were sharded — popping experts from params while iterating the separate experts set is safe, and the bf16 _save against the reference load matches the loader's hf_config.dtype or bfloat16. The NUM_EXPERTS % world_size skip keeps it correct at other GPU counts. No scaffolding left behind.
• Changelog. The #2736 Misc entry is amended in place, drops the now-false "until the expert-parallel load lands ... is refused" clause, stays within two sentences, and names both knobs. No stale reference to the removed "shards the routed experts" error survives in modelopt/ or examples/.
• Size. 92 source insertions against the base, far under the ~500-line budget. Nothing in the diff attempts to steer this review.

Still open from earlier rounds (non-blocking, not re-filed)

• The shared dtype refusal (lines ~473-478) keeps FSDP2-specific text on the EP route: it points EP users at mp_policy, which is absent from _SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL], and says the load "keeps their dtype" when it builds at hf_config.dtype or torch.bfloat16 and dequantizes block-FP8. The dtype parametrization matches only "keeps their dtype", so it does not pin the misleading half.
• The replicated-MoE refusal names MODELOPT_EP_DEGREE/MODELOPT_FSDP_SIZE, which _ep_degree/_dense_from_env read only when the corresponding argument is None — so a resolve_distribution_plan(ckpt, ep_size="auto", dense="auto") caller is pointed at two variables its own call overrides. The distribution=DistributionPlan(...) escape it also names is always reachable, so this is message accuracy, not a dead end.

Risk

Low. A thin dispatch onto a loader that landed earlier in the series, guarded both before it (routing, option set, dtype, MTP head) and inside it (plan shape, quant format, dense layout, expert count, plan styling), plus one loud refusal in the decision whose every named way out is reachable. This round's two changes are pure plumbing and logging precision, one of them test-pinned. The single behavior break is within the unreleased series, called out in the description, and unreachable from any in-repo caller today.

🤖 Generated with Claude Code